#!/usr/bin/python3
# coding=utf-8

# 在macOS, python3.8下测试通过. 需要pip安装一些依赖包, 见import.
# 可以配置个alias: 
# alias z='/path/to/this/jump.py'
# 然后 `z model10` 即可.
# 这个脚本用于safe,hulk等环境机器.

import re, os, sys, time, json
import pexpect, struct, fcntl, termios, signal

user_name = "tianzhipeng-jk" #登录堡垒机的用户名 todo
jumper_server_1 = "jumper.qihoo.net" #堡垒机的ip
jumper_port_1 = 2222 #堡垒机的端口
pin_password = "6位动态码"  #pin码, 拼接完整密码的前半部分 todo
server_user_name = "j-tianzhipeng-jk" #登录服务器的用户名 todo
server_password = "固定服务器密码" #登录服务器的密码 todo

server_2_dict = {
    "data2": "data2.credit.lycc.qihoo.net",
    "data09": "data09.credit.lycc.qihoo.net",
    "cd001v":"cd001v.dataeco.bjyt.qihoo.net",
    "prod02":"prod02.jietiao.bjyt.qihoo.net",
    "datahub01v": "datahub01v.dataeco.bjmd.qihoo.net"
}

def get_terminal_size():
    s = struct.pack("HHHH", 0, 0, 0, 0)
    a = struct.unpack("hhhh", fcntl.ioctl(sys.stdout.fileno(), termios.TIOCGWINSZ, s))
    return a[0], a[1]

def sigwinch_passthrough(sig, data):
    global p
    if not p.closed:
        p.setwinsize(*get_terminal_size())

def strange_sendline(p, str):
    for c in str:
        p.send(c)
    time.sleep(0.1)
    p.sendcontrol('m')

if __name__ == "__main__":
    if len(sys.argv) == 1:
        print("请输入服务器名")
        print(json.dumps(server_2_dict, indent=2))
        sys.exit(0)
    server_2 = sys.argv[1]
    server_name = server_2_dict.get(server_2)
    passwd = input("推推动态码:")
    p = pexpect.spawn(
        "ssh %s@%s -p%d" % (user_name, jumper_server_1, jumper_port_1),
        encoding="utf-8",
        timeout=10,
    )
    p.expect("password:")
    p.sendline(pin_password + passwd)
    p.expect(re.escape("Opt>"))
    strange_sendline(p, server_name)

    p.expect(re.escape("username:"))
    strange_sendline(p, server_user_name)

    rt_no = p.expect(["password\:", "复用SSH连接"])
    if rt_no == 0:
        strange_sendline(p, server_password)

    p.expect(re.escape(server_user_name + "@"))

    # Set the window size the same of current terminal window size
    p.setwinsize(*get_terminal_size())
    # Hook the window change signal so that the pexpect window size change as well
    signal.signal(signal.SIGWINCH, sigwinch_passthrough)
    p.interact()
